import logging
import torch

from trainer import PretrainTrainer, FinetuneTrainer, SampleRunner, ProbGuessRunner
from config.read_properties import read_properties
from config.model_args import ModelArgs

logging.basicConfig(level=logging.INFO)
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)

if __name__ == '__main__':
    args: ModelArgs = read_properties("./config/args.yaml")
    device = torch.device(args.cuda_devices)

    if args.func == "pretrain":
        trainer = PretrainTrainer(args, device)
        trainer.run()
    elif args.func == "finetune":
        trainer = FinetuneTrainer(args, device)
        trainer.run()
    elif args.func == "sample":
        trainer = SampleRunner(args, device)
        trainer.run()
    elif args.func == "prob_guess":
        trainer = ProbGuessRunner(args, device)
        trainer.run()
    else:
        trainer = ProbGuessRunner(args, device)
        trainer.run()
